#include <stdio.h>
#include "mat.h"
#include <opencv2/opencv.hpp>

namespace ncnn {
    static double get_current_time();
}

static ncnn::Mat from_rgb_c(const unsigned char* rgb, int w, int h);
static ncnn::Mat from_rgb_neon_intrinsics(const unsigned char* rgb, int w, int h);

static ncnn::Mat from_rgb2bgr_c(const unsigned char* rgb, int w, int h);
static ncnn::Mat from_rgb2bgr_neon_intrinsics(const unsigned char* rgb, int w, int h);

static void to_rgb_c(const ncnn::Mat& m, unsigned char* rgb);
static void to_rgb_neon_intrinsics(const ncnn::Mat& m, unsigned char* rgb);

namespace ncnn {
#ifdef _WIN32
    double get_current_time()
{
    LARGE_INTEGER freq;
    LARGE_INTEGER pc;
    QueryPerformanceFrequency(&freq);
    QueryPerformanceCounter(&pc);

    return pc.QuadPart * 1000.0 / freq.QuadPart;
}
#else  // _WIN32

    double get_current_time() {
        struct timeval tv;
        gettimeofday(&tv, NULL);

        return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
    }

#endif // _WIN32
}

// C版本的from_rgb实现
ncnn::Mat from_rgb_c(const unsigned char* rgb, int w, int h) {
    ncnn::Mat m(w, h, 3); // 创建3通道图像
    if (m.empty()) {
        return m;
    }

    // 创建3个指针，分别指向刚创建图像的3个通道
    // 注意：源buffer的每个bgr是连续的，而目标buffer的bgr是分开的
    // 也就是：ncnn::Mat是CHW顺序，而cv::Mat是HWC顺序
    // 因此，from_rgb函数，不能简单的用memcpy实现
    // 但如果用NHWC顺序做CNN推理，则from_rgb可用memcpy甚至浅拷贝实现
    float* ptr0 = m.channel(0);
    float* ptr1 = m.channel(1);
    float* ptr2 = m.channel(2);

    int size = w * h;
    int remain = size;
    for (int i=0; i<remain; i++) {
        *ptr0 = rgb[0];
        *ptr1 = rgb[1];
        *ptr2 = rgb[2];

        rgb += 3;
        ptr0++;
        ptr1++;
        ptr2++;
    }
    return m;
}

// neon intrinsics版本的from_rgb
ncnn::Mat from_rgb_neon_intrinsics(const unsigned char* rgb, int w, int h)
{
    ncnn::Mat m(w, h, 3);
    if (m.empty()) {
        return m;
    }

    float* ptr0 = m.channel(0);
    float* ptr1 = m.channel(1);
    float* ptr2 = m.channel(2);

    int size = w * h;

    int nn = size >> 3; // size / 8
    int remain = size - (nn<<3); // size - ((size/8)*8)

    for (; nn>0; nn--) {
        // vld3_u8: 按交叉顺序将内存的数据装入3个neon寄存器(d寄存器）
        // 其中每个d寄存器的每个通道(lane)是8个bit，每个d寄存器是64位，也就是8个lane
        // 通道(lane)的意思，是把寄存器看做是基本长度(size)的向量，每个向量元素是一个lane
        // 内存第1个数据放入第1个d寄存器的第1个通道，
        // 内存第2个数据放入第2个d寄存器的第1个通道,
        // 内存第3个数据放入第3个d寄存器的第1个通道,
        // 内存第4个数据放入第1个d寄存器的第2个通道,
        // 内存第5个数据放入第2个d寄存器的第2个通道,
        // 内存第6个数据放入第3个d寄存器的第2个通道,
        // ...
        // 内存第22个数据放入第1个d寄存器的第8个通道,
        // 内存第23个数据放入第2个d寄存器的第8个通道,
        // 内存第24个数据放入第3个d寄存器的第8个通道,
        uint8x8x3_t _rgb = vld3_u8(rgb);

        // 将vector的元素bit位扩大到原来的两倍，元素值不变。
        // 目的是什么？防止溢出吗？其实是为了u8=>float32的类型转换
        // 通常，cnn的输入是float类型，手动u8=>float效率并不高
        // 这里是把u8=>float用neon intrinsic实现了
        uint16x8_t _r16 = vmovl_u8(_rgb.val[0]);
        uint16x8_t _g16 = vmovl_u8(_rgb.val[1]);
        uint16x8_t _b16 = vmovl_u8(_rgb.val[2]);

        // vget_low_u16(_r16)是取_r16的低一半（64位)
        float32x4_t _rlow = vcvtq_f32_u32(vmovl_u16(vget_low_u16(_r16)));
        // vget_high_u16(_r16)是取_r16的高一半（64位）
        float32x4_t _rhigh = vcvtq_f32_u32(vmovl_u16(vget_high_u16(_r16)));

        float32x4_t _glow = vcvtq_f32_u32(vmovl_u16(vget_low_u16(_g16)));
        float32x4_t _ghigh = vcvtq_f32_u32(vmovl_u16(vget_high_u16(_g16)));

        float32x4_t _blow = vcvtq_f32_u32(vmovl_u16(vget_low_u16(_b16)));
        float32x4_t _bhigh = vcvtq_f32_u32(vmovl_u16(vget_high_u16(_b16)));

        vst1q_f32(ptr0, _rlow);
        vst1q_f32(ptr0+4, _rhigh);

        vst1q_f32(ptr1, _glow);
        vst1q_f32(ptr1+4, _ghigh);

        vst1q_f32(ptr2, _blow);
        vst1q_f32(ptr2+4, _bhigh);

        rgb += 3*8;
        ptr0 += 8;
        ptr1 += 8;
        ptr2 += 8;
    }

    for (; remain>0; remain--) {
        *ptr0 = rgb[0];
        *ptr1 = rgb[1];
        *ptr2 = rgb[2];

        rgb += 3;
        ptr0++;
        ptr1++;
        ptr2++;
    }

    return m;
}


/**
 * C版本的from_rgb2bgr实现
 * 维度顺序：HWC => CHW
 * 数据类型：uchar => float
 * 通道顺序：rgb => bgr
 * @param rgb 图像raw buffer，HWC顺序
 * @param w 图像宽度
 * @param h 图像高度
 * @return Mat对象，CHW顺序
 */
ncnn::Mat from_rgb2bgr_c(const unsigned char* rgb, int w, int h) {
    ncnn::Mat m(w, h, 3); // 创建3通道图像
    if (m.empty()) {
        return m;
    }

    // 创建数据指针，存放的是Mat的rgb的每个通道起始地址
    float* ptr0 = m.channel(0);
    float* ptr1 = m.channel(1);
    float* ptr2 = m.channel(2);

    int size = w * h;
    int remain = size;
    for (; remain>0; remain--) {
        // 交错赋值，实现rgb2bgr，或者bgr2rgb功能
        *ptr0 = rgb[2];
        *ptr1 = rgb[1];
        *ptr2 = rgb[1];

        rgb += 3;
        ptr0++;
        ptr1++;
        ptr2++;
    }
    return m;
}


/**
 * neon intrinsics版本的from_rgb2bgr实现
 * 维度顺序：HWC => CHW
 * 数据类型：uchar => float
 * 通道顺序：rgb => bgr
 * @param rgb 图像raw buffer，HWC顺序
 * @param w 图像宽度
 * @param h 图像高度
 * @return Mat对象，CHW顺序
 */
ncnn::Mat from_rgb2bgr_neon_intrinsics(const unsigned char* rgb, int w, int h) {
    ncnn::Mat m(w, h, 3); // 创建3通道图像
    if (m.empty()) {
        return m;
    }

    // 创建数据指针，存放的是Mat的rgb的每个通道起始地址
    float* ptr0 = m.channel(0);
    float* ptr1 = m.channel(1);
    float* ptr2 = m.channel(2);

    int size = w * h;
    int nn = size >> 3;
    int remain = size - (nn<<3);
    for(; nn>0; nn--) {
        uint8x8x3_t _rgb = vld3_u8(rgb);
        uint16x8_t _r16 = vmovl_u8(_rgb.val[0]);
        uint16x8_t _g16 = vmovl_u8(_rgb.val[1]);
        uint16x8_t _b16 = vmovl_u8(_rgb.val[2]);

        float32x4_t _rlow = vcvtq_f32_u32(vmovl_u16(vget_low_u16(_r16)));
        float32x4_t _rhigh = vcvtq_f32_u32(vmovl_u16(vget_high_u16(_g16)));

        float32x4_t _glow = vcvtq_f32_u32(vmovl_u16(vget_low_u16(_g16)));
        float32x4_t _ghigh = vcvtq_f32_u32(vmovl_u16(vget_high_u16(_g16)));

        float32x4_t _blow = vcvtq_f32_u32(vmovl_u16(vget_low_u16(_b16)));
        float32x4_t _bhigh = vcvtq_f32_u32(vmovl_u16(vget_high_u16(_b16)));

        vst1q_f32(ptr2, _rlow);
        vst1q_f32(ptr2+4, _rhigh);

        vst1q_f32(ptr1, _glow);
        vst1q_f32(ptr1+4, _ghigh);

        vst1q_f32(ptr2, _blow);
        vst1q_f32(ptr2+4, _bhigh);

        rgb += 3*8;
        ptr0 += 8;
        ptr1 += 8;
        ptr2 += 8;
    }

    for (; remain>0; remain--) {
        // 交错赋值，实现rgb2bgr，或者bgr2rgb功能
        *ptr0 = rgb[2];
        *ptr1 = rgb[1];
        *ptr2 = rgb[1];

        rgb += 3;
        ptr0++;
        ptr1++;
        ptr2++;
    }
    return m;
}


/**
 * Mat转buffer
 * 维度：HWC => CHW
 * 类型：float => uchar
 * @param m
 * @param rgb
 */
static void to_rgb_c(const ncnn::Mat& m, unsigned char* rgb)
{
    const float* ptr0 = m.channel(0);
    const float* ptr1 = m.channel(1);
    const float* ptr2 = m.channel(2);

    int size = m.w * m.h;

//#define SATURATE_CAST_UCHAR(X) (unsigned char)std::min(std::max((int)(X), 0), 255);
#define SATURATE_CAST_UCHAR(X) (unsigned char)(X)
    int remain = size;

    for (; remain>0; remain--)
    {
        rgb[0] = SATURATE_CAST_UCHAR(*ptr0);
        rgb[1] = SATURATE_CAST_UCHAR(*ptr1);
        rgb[2] = SATURATE_CAST_UCHAR(*ptr2);

        rgb += 3;
        ptr0++;
        ptr1++;
        ptr2++;
    }

#undef SATURATE_CAST_UCHAR
}

// 改动1： vld1q_dup_f32 => vld1q_f32
void to_rgb_neon_intrinsics(const ncnn::Mat& m, unsigned char* rgb)
{
    const float* ptr0 = m.channel(0);
    const float* ptr1 = m.channel(1);
    const float* ptr2 = m.channel(2);

    int size = m.w * m.h;
    int nn = size >> 3;
    int remain = size - (nn<<3);
    for (; nn>0; nn--) {
        // r
        float32x4_t _rlow = vld1q_f32(ptr0);
        float32x4_t _rhigh = vld1q_f32(ptr0+4);

        float32x4_t _glow = vld1q_f32(ptr1);
        float32x4_t _ghigh = vld1q_f32(ptr1+4);

        float32x4_t _blow = vld1q_f32(ptr2);
        float32x4_t _bhigh = vld1q_f32(ptr2+4);

        uint16x8_t _r16 = vcombine_u16(vmovn_u32(vcvtq_u32_f32(_rlow)), vmovn_u32(vcvtq_u32_f32(_rhigh)));
        uint16x8_t _g16 = vcombine_u16(vmovn_u32(vcvtq_u32_f32(_glow)), vmovn_u32(vcvtq_u32_f32(_ghigh)));
        uint16x8_t _b16 = vcombine_u16(vmovn_u32(vcvtq_u32_f32(_blow)), vmovn_u32(vcvtq_u32_f32(_bhigh)));

        uint8x8x3_t _rgb;
        _rgb.val[0] = vmovn_u16(_r16);
        _rgb.val[1] = vmovn_u16(_g16);
        _rgb.val[2] = vmovn_u16(_b16);
        vst3_u8(rgb, _rgb);

        rgb += 3*8;
        ptr0 += 8;
        ptr1 += 8;
        ptr2 += 8;
    }
    for (; remain>0; remain--) {
#define SATURATE_CAST_UCHAR(X) (unsigned char)::std::min(::std::max((int)(X), 0), 255);
        rgb[0] = SATURATE_CAST_UCHAR(*ptr0);
        rgb[1] = SATURATE_CAST_UCHAR(*ptr1);
        rgb[2] = SATURATE_CAST_UCHAR(*ptr2);
        rgb += 3;
        ptr0++;
        ptr1++;
        ptr2++;
#undef SATURATE_CAST_UCHAR
    }
}



// 输入图像尺寸 width=3880, size=5184
// 重复10次
// from_rgb_c耗时 1320 ms
// from_rgb_neon_intrinsics耗时 1230 ms
int main() {

    printf("hello ncnn!\n");
    //std::string filename = "mingren.jpg";
    std::string filename = "FaceEdit_20201231203904.png";
    cv::Mat image = cv::imread(filename);
    cv::Size size = image.size();
    // int loop = 10;
    // double t_start = ncnn::get_current_time();
    // for (int i=0; i<loop; i++) {
    //     ncnn::Mat mat1 = from_rgb_c(image.data, size.width, size.height);
    // }
    // double t_cost = ncnn::get_current_time() - t_start;

    // t_start = ncnn::get_current_time();
    // for (int i=0; i<loop; i++) {
    //     ncnn::Mat mat2 = from_rgb_neon_intrinsics(image.data, size.width, size.height);
    // }
    // double t_cost2 = ncnn::get_current_time() - t_start;

    // printf("image size: width=%d, size=%d\n", size.width, size.height);
    // printf("create_mat_from_rgb, time cost %lf ms\n", t_cost);
    // printf("create_mat_from_rgb, time cost %lf ms\n", t_cost2);

    ncnn::Mat mat = from_rgb_neon_intrinsics(image.data, size.width, size.height);
    //ncnn::Mat mat = from_rgb_c(image.data, size.width, size.height);

    cv::Mat result(size, CV_8UC3);
    to_rgb_neon_intrinsics(mat, result.data);
    //to_rgb_neon_intrinsics(mat, result.data);
    cv::imwrite("result.png", result);


    return 0;
}